Amazon SageMakerにおけるRecordIO形式のデータの作成と読み込み

どうも、DA事業本部の大澤です

SageMakerの組み込みアルゴリズムではデータ形式として、RecordIOに対応していることがよくあります。 今回はSageMakerで使われるRecordIO形式のデータの作成と読み込みについて紹介します。

バージョン情報

今回紹介する内容はSageMaker Python SDKに強く依存するため、バージョン情報を書いておきます。

• SageMaker Python SDK 1.36.4

ライブラリの読み込み

まずは必要なライブラリを読み込みます。

import numpy as np
from scipy import sparse
import sagemaker.amazon.common as smac
import tempfile
from sagemaker.amazon.record_pb2 import Record

ndarrayとRecordIO

NumPyの配列であるndarrayからRecordIO形式のファイルに変換し、RecordIO形式のファイルを読み込んで値を確認します。

array_data = np.array([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]])
labels = np.array([1.0, 0.0])
with tempfile.TemporaryFile() as f:
    # ndarrayからRecordIOに変換
    smac.write_numpy_to_dense_tensor(f, array_data, labels)

    f.seek(0)

    # RecordIOを読み込む
    records = smac.read_recordio(f)

    # 行ごとにデータを取り出す
    for record_string in records:
        print('record_string', record_string)
        record = Record()
        record.ParseFromString(record_string)
        print('record', record)
        print('values', record.features['values'].float64_tensor.values)
        print('label', record.label['values'].float64_tensor.values)

出力

record_string b'\n\x1c\n\x06values\x12\x12\x1a\x10\n\x08\x00\x00\x00\x00\x00\x00\xf0?\x12\x01\x01\x1a\x01\x03\x12\x16\n\x06values\x12\x0c\x1a\n\n\x08\x00\x00\x00\x00\x00\x00\x00\x00'
record features {
key: "values"
value {
float64_tensor {
values: 1.0
keys: 1
shape: 3
}
}
}
label {
key: "values"
value {
float64_tensor {
values: 0.0
}
}
}

keys [1]
values [1.0]
shape [3]
label [0.0]
record_string b'\n%\n\x06values\x12\x1b\x1a\x19\n\x10\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?\x12\x02\x00\x02\x1a\x01\x03\x12\x16\n\x06values\x12\x0c\x1a\n\n\x08\x00\x00\x00\x00\x00\x00\xf0?'
record features {
key: "values"
value {
float64_tensor {
values: 1.0
values: 1.0
keys: 0
keys: 2
shape: 3
}
}
}
label {
key: "values"
value {
float64_tensor {
values: 1.0
}
}
}

今回はfloat64形式のデータなので、record.features['values'].float64_tensor.valuesで値を参照しています。float32であればrecord.features['values'].float32_tensor.valuesといった形で参照方法が変化します。

疎行列(scipy.sparse)とRecordIO

疎行列であるcoo_matrixからRecordIO形式のファイルに変換し、RecordIO形式のファイルを読み込んで値を確認します。 今回はcoo形式の疎行列で試していますが、cscやcsr、lil形式の疎行列でも同様に変換可能です。

array_data = [[0.0, 1.0, 0.0], [1.0, 0.0, 1.0]]
labels = np.array([0.0, 1.0])
sparse_array = sparse.coo_matrix(np.array(array_data))
with tempfile.TemporaryFile() as f:
    # 疎行列からRecordIOに変換
    smac.write_spmatrix_to_sparse_tensor(f, sparse_array, labels)

    f.seek(0)

    # RecordIOを読み込む
    records = smac.read_recordio(f)

    # 行ごとにデータを取り出す
    for record_string in records:
        print('record_string', record_string)
        record = Record()
        record.ParseFromString(record_string)
        print('record', record)
        print('keys', record.features['values'].float64_tensor.keys)
        print('values', record.features['values'].float64_tensor.values)
        print('shape', record.features['values'].float64_tensor.shape)
        print('label', record.label['values'].float64_tensor.values)

出力

record_string b'\n\x1c\n\x06values\x12\x12\x1a\x10\n\x08\x00\x00\x00\x00\x00\x00\xf0?\x12\x01\x01\x1a\x01\x03\x12\x16\n\x06values\x12\x0c\x1a\n\n\x08\x00\x00\x00\x00\x00\x00\x00\x00'
record features {
key: "values"
value {
float64_tensor {
values: 1.0
keys: 1
shape: 3
}
}
}
label {
key: "values"
value {
float64_tensor {
values: 0.0
}
}
}

keys [1]
values [1.0]
shape [3]
label [0.0]
record_string b'\n%\n\x06values\x12\x1b\x1a\x19\n\x10\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?\x12\x02\x00\x02\x1a\x01\x03\x12\x16\n\x06values\x12\x0c\x1a\n\n\x08\x00\x00\x00\x00\x00\x00\xf0?'
record features {
key: "values"
value {
float64_tensor {
values: 1.0
values: 1.0
keys: 0
keys: 2
shape: 3
}
}
}
label {
key: "values"
value {
float64_tensor {
values: 1.0
}
}
}

keys [0, 2]
values [1.0, 1.0]
shape [3]
label [1.0]

RecordIOに変換する際に、write_spmatrix_to_sparse_tensorを使ってる点が、ndarrayの場合と異なります。 読み込み方はndarrayの場合と基本的に同じですが、疎行列から変換した場合はkeysとshapeというデータが追加されています。

疎行列なので、次のような特定の列が値を持ち、それ以外が0みたいな行列を扱うケースが多いかと思います。 そういった場合でもRecordIOでは必要なデータのみを保持するようになっています。 また、今回はラベル無しで試してみます。

array_data = [[0.0, 1.0, 0.0], [1.0, 0.0, 1.0]]
sparse_array = sparse.coo_matrix(np.array(array_data))
with tempfile.TemporaryFile() as f:
    smac.write_spmatrix_to_sparse_tensor(f, sparse_array)
    f.seek(0)
    records = smac.read_recordio(f)
    for record_string in records:
        print('record_string', record_string)
        record = Record()
        record.ParseFromString(record_string)
        print('record', record)
        print('keys', record.features['values'].float64_tensor.keys)
        print('values', record.features['values'].float64_tensor.values)
        print('shape', record.features['values'].float64_tensor.shape)

出力

record_string b'\n\x1c\n\x06values\x12\x12\x1a\x10\n\x08\x00\x00\x00\x00\x00\x00\xf0?\x12\x01\x01\x1a\x01\x03'
record features {
key: "values"
value {
float64_tensor {
values: 1.0
keys: 1
shape: 3
}
}
}

keys [1]
values [1.0]
shape [3]
record_string b'\n%\n\x06values\x12\x1b\x1a\x19\n\x10\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?\x12\x02\x00\x02\x1a\x01\x03'
record features {
key: "values"
value {
float64_tensor {
values: 1.0
values: 1.0
keys: 0
keys: 2
shape: 3
}
}
}

keys [0, 2]
values [1.0, 1.0]
shape [3]

必要な値のみが格納されていることがわかります。今回は6要素しかありませんが、これが数百万数千万もしくはそれ以上のデータを扱う場合には省メモリ効果がかなり大きくなります。

さいごに

SageMakerの組み込みアルゴリズムでよく使われるRecordIO形式のデータの作成方法と読み込み方法について紹介しました。 SageMakerのサンプルノートブックに従ってRecordIO形式のデータを作ったけども、これどうやって読み込むんやろか...みたいな時等に試していただければと思います。

参考

• sagemaker-python-sdk/test_common.py at master · aws/sagemaker-python-sdk
• sagemaker-python-sdk/common.py at master · aws/sagemaker-python-sdk